[MLIR][Python] Support eliding large resource strings in PassManager #149187

akshaykhadse · 2025-07-16T21:21:19Z

Introduces a large_resource_limit parameter across Python bindings, enabling the eliding of resource strings exceeding a specified character limit during IR printing.
To maintain backward compatibilty, when using operation.print() API, if large_resource_limit is None and the large_elements_limit is set, the later will be used to elide the resource string as well. This change was introduced by [mlir] Python: Extend print large elements limit to resources #125738.
For printing using pass manager, the large_resource_limit and large_elements_limit are completely independent of each other.

github-actions · 2025-07-16T21:23:39Z

✅ With the latest revision this PR passed the Python code formatter.

llvmbot · 2025-07-16T22:41:17Z

@llvm/pr-subscribers-mlir

Author: Akshay Khadse (akshaykhadse)

Changes

Introduces a large_resource_limit parameter across Python bindings, enabling the elision of resource strings exceeding a specified character limit during IR printing.
Reverts the behavior where large_elements_limit would also control the resource string printing from operation's print method when using the Python binding. This is because the eliding of resource string takes the number of characters as input whereas the elision of elements takes number of elements as input. Both are separate and should be respected in Python bindings as well.

Full diff: https://github.com/llvm/llvm-project/pull/149187.diff

7 Files Affected:

(modified) mlir/lib/Bindings/Python/IRCore.cpp (+14-6)
(modified) mlir/lib/Bindings/Python/IRModule.h (+7-7)
(modified) mlir/lib/Bindings/Python/Pass.cpp (+6-1)
(modified) mlir/python/mlir/_mlir_libs/_mlir/ir.pyi (+5)
(modified) mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi (+1)
(modified) mlir/test/python/ir/operation.py (+9)
(modified) mlir/test/python/pass_manager.py (+57)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index d961482885300..c90b866e67ea0 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -97,6 +97,8 @@ static const char kOperationPrintDocstring[] =
   binary: Whether to write bytes (True) or str (False). Defaults to False.
   large_elements_limit: Whether to elide elements attributes above this
     number of elements. Defaults to None (no limit).
+  large_resource_limit: Whether to elide resource attributes above this
+    number of characters. Defaults to None (no limit).
   enable_debug_info: Whether to print debug/location information. Defaults
     to False.
   pretty_debug_info: Whether to format debug information for easier reading
@@ -1303,6 +1305,7 @@ void PyOperation::checkValid() const {
 }
 
 void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
+                            std::optional<int64_t> largeResourceLimit,
                             bool enableDebugInfo, bool prettyDebugInfo,
                             bool printGenericOpForm, bool useLocalScope,
                             bool useNameLocAsPrefix, bool assumeVerified,
@@ -1314,10 +1317,10 @@ void PyOperationBase::print(std::optional<int64_t> largeElementsLimit,
     fileObject = nb::module_::import_("sys").attr("stdout");
 
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
-  if (largeElementsLimit) {
+  if (largeElementsLimit)
     mlirOpPrintingFlagsElideLargeElementsAttrs(flags, *largeElementsLimit);
-    mlirOpPrintingFlagsElideLargeResourceString(flags, *largeElementsLimit);
-  }
+  if (largeResourceLimit)
+    mlirOpPrintingFlagsElideLargeResourceString(flags, *largeResourceLimit);
   if (enableDebugInfo)
     mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
                                        /*prettyForm=*/prettyDebugInfo);
@@ -1405,6 +1408,7 @@ void PyOperationBase::walk(
 
 nb::object PyOperationBase::getAsm(bool binary,
                                    std::optional<int64_t> largeElementsLimit,
+                                   std::optional<int64_t> largeResourceLimit,
                                    bool enableDebugInfo, bool prettyDebugInfo,
                                    bool printGenericOpForm, bool useLocalScope,
                                    bool useNameLocAsPrefix, bool assumeVerified,
@@ -1416,6 +1420,7 @@ nb::object PyOperationBase::getAsm(bool binary,
     fileObject = nb::module_::import_("io").attr("StringIO")();
   }
   print(/*largeElementsLimit=*/largeElementsLimit,
+        /*largeResourceLimit=*/largeResourceLimit,
         /*enableDebugInfo=*/enableDebugInfo,
         /*prettyDebugInfo=*/prettyDebugInfo,
         /*printGenericOpForm=*/printGenericOpForm,
@@ -3348,6 +3353,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           [](PyOperationBase &self) {
             return self.getAsm(/*binary=*/false,
                                /*largeElementsLimit=*/std::nullopt,
+                               /*largeResourceLimit=*/std::nullopt,
                                /*enableDebugInfo=*/false,
                                /*prettyDebugInfo=*/false,
                                /*printGenericOpForm=*/false,
@@ -3363,11 +3369,12 @@ void mlir::python::populateIRCore(nb::module_ &m) {
            nb::arg("state"), nb::arg("file").none() = nb::none(),
            nb::arg("binary") = false, kOperationPrintStateDocstring)
       .def("print",
-           nb::overload_cast<std::optional<int64_t>, bool, bool, bool, bool,
-                             bool, bool, nb::object, bool, bool>(
-               &PyOperationBase::print),
+           nb::overload_cast<std::optional<int64_t>, std::optional<int64_t>,
+                             bool, bool, bool, bool, bool, bool, nb::object,
+                             bool, bool>(&PyOperationBase::print),
            // Careful: Lots of arguments must match up with print method.
            nb::arg("large_elements_limit").none() = nb::none(),
+           nb::arg("large_resource_limit").none() = nb::none(),
            nb::arg("enable_debug_info") = false,
            nb::arg("pretty_debug_info") = false,
            nb::arg("print_generic_op_form") = false,
@@ -3383,6 +3390,7 @@ void mlir::python::populateIRCore(nb::module_ &m) {
            // Careful: Lots of arguments must match up with get_asm method.
            nb::arg("binary") = false,
            nb::arg("large_elements_limit").none() = nb::none(),
+           nb::arg("large_resource_limit").none() = nb::none(),
            nb::arg("enable_debug_info") = false,
            nb::arg("pretty_debug_info") = false,
            nb::arg("print_generic_op_form") = false,
diff --git a/mlir/lib/Bindings/Python/IRModule.h b/mlir/lib/Bindings/Python/IRModule.h
index 9befcce725bb7..0fdd2d1a7eff6 100644
--- a/mlir/lib/Bindings/Python/IRModule.h
+++ b/mlir/lib/Bindings/Python/IRModule.h
@@ -599,18 +599,18 @@ class PyOperationBase {
 public:
   virtual ~PyOperationBase() = default;
   /// Implements the bound 'print' method and helps with others.
-  void print(std::optional<int64_t> largeElementsLimit, bool enableDebugInfo,
+  void print(std::optional<int64_t> largeElementsLimit,
+             std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
              bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
              bool useNameLocAsPrefix, bool assumeVerified,
              nanobind::object fileObject, bool binary, bool skipRegions);
   void print(PyAsmState &state, nanobind::object fileObject, bool binary);
 
-  nanobind::object getAsm(bool binary,
-                          std::optional<int64_t> largeElementsLimit,
-                          bool enableDebugInfo, bool prettyDebugInfo,
-                          bool printGenericOpForm, bool useLocalScope,
-                          bool useNameLocAsPrefix, bool assumeVerified,
-                          bool skipRegions);
+  nanobind::object
+  getAsm(bool binary, std::optional<int64_t> largeElementsLimit,
+         std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
+         bool prettyDebugInfo, bool printGenericOpForm, bool useLocalScope,
+         bool useNameLocAsPrefix, bool assumeVerified, bool skipRegions);
 
   // Implement the bound 'writeBytecode' method.
   void writeBytecode(const nanobind::object &fileObject,
diff --git a/mlir/lib/Bindings/Python/Pass.cpp b/mlir/lib/Bindings/Python/Pass.cpp
index 8d84864b9db4d..830e106499fc1 100644
--- a/mlir/lib/Bindings/Python/Pass.cpp
+++ b/mlir/lib/Bindings/Python/Pass.cpp
@@ -78,12 +78,16 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
           [](PyPassManager &passManager, bool printBeforeAll,
              bool printAfterAll, bool printModuleScope, bool printAfterChange,
              bool printAfterFailure, std::optional<int64_t> largeElementsLimit,
-             bool enableDebugInfo, bool printGenericOpForm,
+             std::optional<int64_t> largeResourceLimit, bool enableDebugInfo,
+             bool printGenericOpForm,
              std::optional<std::string> optionalTreePrintingPath) {
             MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
             if (largeElementsLimit)
               mlirOpPrintingFlagsElideLargeElementsAttrs(flags,
                                                          *largeElementsLimit);
+            if (largeResourceLimit)
+              mlirOpPrintingFlagsElideLargeResourceString(flags,
+                                                          *largeResourceLimit);
             if (enableDebugInfo)
               mlirOpPrintingFlagsEnableDebugInfo(flags, /*enable=*/true,
                                                  /*prettyForm=*/false);
@@ -103,6 +107,7 @@ void mlir::python::populatePassManagerSubmodule(nb::module_ &m) {
           "print_module_scope"_a = false, "print_after_change"_a = false,
           "print_after_failure"_a = false,
           "large_elements_limit"_a.none() = nb::none(),
+          "large_resource_limit"_a.none() = nb::none(),
           "enable_debug_info"_a = false, "print_generic_op_form"_a = false,
           "tree_printing_dir_path"_a.none() = nb::none(),
           "Enable IR printing, default as mlir-print-ir-after-all.")
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
index ed476da28d6be..d0212c7a1e126 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/ir.pyi
@@ -200,6 +200,7 @@ class _OperationBase:
     def get_asm(
         binary: Literal[True],
         large_elements_limit: int | None = None,
+        large_resource_limit: int | None = None,
         enable_debug_info: bool = False,
         pretty_debug_info: bool = False,
         print_generic_op_form: bool = False,
@@ -212,6 +213,7 @@ class _OperationBase:
         self,
         binary: bool = False,
         large_elements_limit: int | None = None,
+        large_resource_limit: int | None = None,
         enable_debug_info: bool = False,
         pretty_debug_info: bool = False,
         print_generic_op_form: bool = False,
@@ -253,6 +255,7 @@ class _OperationBase:
     def print(
         self,
         large_elements_limit: int | None = None,
+        large_resource_limit: int | None = None,
         enable_debug_info: bool = False,
         pretty_debug_info: bool = False,
         print_generic_op_form: bool = False,
@@ -270,6 +273,8 @@ class _OperationBase:
           binary: Whether to write bytes (True) or str (False). Defaults to False.
           large_elements_limit: Whether to elide elements attributes above this
             number of elements. Defaults to None (no limit).
+          large_resource_limit: Whether to elide resource strings above this
+            number of characters. Defaults to None (no limit).
           enable_debug_info: Whether to print debug/location information. Defaults
             to False.
           pretty_debug_info: Whether to format debug information for easier reading
diff --git a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
index 0d2eaffe16d3e..1010daddae2aa 100644
--- a/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi
@@ -23,6 +23,7 @@ class PassManager:
         print_after_change: bool = False,
         print_after_failure: bool = False,
         large_elements_limit: int | None = None,
+        large_resource_limit: int | None = None,
         enable_debug_info: bool = False,
         print_generic_op_form: bool = False,
         tree_printing_dir_path: str | None = None,
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index b08fe98397fbc..ede1571f940f6 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -686,6 +686,15 @@ def testOperationPrint():
         skip_regions=True,
     )
 
+    # Test print with large_resource_limit.
+    # CHECK: func.func @f1(%arg0: i32) -> i32
+    # CHECK-NOT: resource1: "0x08
+    module.operation.print(large_resource_limit=2)
+
+    # Test large_elements_limit has no effect on resource string
+    # CHECK: func.func @f1(%arg0: i32) -> i32
+    # CHECK: resource1: "0x08
+    module.operation.print(large_elements_limit=2)
 
 # CHECK-LABEL: TEST: testKnownOpView
 @run
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index 85d2eb304882e..4c992a3ab2e45 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -363,6 +363,63 @@ def testPrintIrLargeLimitElements():
         pm.run(module)
 
 
+# CHECK-LABEL: TEST: testPrintIrLargeResourceLimit
+@run
+def testPrintIrLargeResourceLimit():
+    with Context() as ctx:
+        module = ModuleOp.parse(
+            """
+          module {
+            func.func @main() -> tensor<3xi64> {
+              %0 = arith.constant dense_resource<blob1> : tensor<3xi64>
+              return %0 : tensor<3xi64>
+            }
+          }
+          {-#
+            dialect_resources: {
+              builtin: {
+                blob1: "0x010000000000000002000000000000000300000000000000"
+              }
+            }
+          #-}
+        """
+        )
+        pm = PassManager.parse("builtin.module(canonicalize)")
+        ctx.enable_multithreading(False)
+        pm.enable_ir_printing(large_resource_limit=4)
+        # CHECK-NOT: blob1: "0x01
+        pm.run(module)
+
+
+# CHECK-LABEL: TEST: testPrintIrLargeResourceLimitVsElementsLimit
+@run
+def testPrintIrLargeResourceLimitVsElementsLimit():
+    """Test that large_elements_limit does not affect the printing of resources."""
+    with Context() as ctx:
+        module = ModuleOp.parse(
+            """
+          module {
+            func.func @main() -> tensor<3xi64> {
+              %0 = arith.constant dense_resource<blob1> : tensor<3xi64>
+              return %0 : tensor<3xi64>
+            }
+          }
+          {-#
+            dialect_resources: {
+              builtin: {
+                blob1: "0x010000000000000002000000000000000300000000000000"
+              }
+            }
+          #-}
+        """
+        )
+        pm = PassManager.parse("builtin.module(canonicalize)")
+        ctx.enable_multithreading(False)
+        pm.enable_ir_printing(large_elements_limit=1)
+        # CHECK: blob1: "0x01
+        pm.run(module)
+
+
 # CHECK-LABEL: TEST: testPrintIrTree
 @run
 def testPrintIrTree():

mlir/test/python/pass_manager.py

makslevental

Looks good to me but this is a "breaking" change - can you take a look and see if anyone is using the previous API to handle resource strings (you can try https://sourcegraph.com/search). If you don't find anything it's probably safe to change but you'll need to make a post on the discourse about it (like this one https://discourse.llvm.org/t/psa-changes-in-mlir-python-type-casters/87383).

github-actions · 2025-07-17T02:14:00Z

✅ With the latest revision this PR passed the C/C++ code formatter.

- Introduces a `large_resource_limit` parameter across Python bindings, enabling the eliding of resource strings exceeding a specified character limit during IR printing. - To maintain backward compatibilty, when using `operation.print()` API, if `large_resource_limit` is None and the `large_elements_limit` is set, the later will be used to elide the resource string as well. This change was introduced by llvm#125738. - For printing using pass manager, the `large_resource_limit` and `large_elements_limit` are completely independent of each other.

makslevental · 2025-07-17T04:34:49Z

Cool since this is backwards compatible we can merge this - let me know when you're ready and I'll merge it.

akshaykhadse force-pushed the fix-resource-string-printing branch 2 times, most recently from a904185 to 658a9ef Compare July 16, 2025 21:47

akshaykhadse marked this pull request as ready for review July 16, 2025 22:40

akshaykhadse requested review from ftynse, makslevental and stellaraccident as code owners July 16, 2025 22:40

llvmbot added mlir:python MLIR Python bindings mlir labels Jul 16, 2025

makslevental reviewed Jul 16, 2025

View reviewed changes

mlir/test/python/pass_manager.py Outdated Show resolved Hide resolved

makslevental approved these changes Jul 17, 2025

View reviewed changes

akshaykhadse force-pushed the fix-resource-string-printing branch from 658a9ef to 60479f3 Compare July 17, 2025 02:11

akshaykhadse force-pushed the fix-resource-string-printing branch from 60479f3 to 6323eeb Compare July 17, 2025 02:25

akshaykhadse self-assigned this Jul 17, 2025

akshaykhadse changed the title ~~Support eliding large resource strings in Python PassManager~~ [MLIR][Python] Support eliding large resource strings in Python PassManager Jul 17, 2025

akshaykhadse changed the title ~~[MLIR][Python] Support eliding large resource strings in Python PassManager~~ [MLIR][Python] Support eliding large resource strings in PassManager Jul 17, 2025

makslevental merged commit e4a3541 into llvm:main Jul 17, 2025
9 checks passed

tru mentioned this pull request Jul 17, 2025

Use Parallel xz for test-suite sources. #149389

Merged

This was referenced Jul 23, 2025

test abhinavgaba/llvm-project#2

Closed

Add dataFence plugin interface abhinavgaba/llvm-project#3

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][Python] Support eliding large resource strings in PassManager #149187

[MLIR][Python] Support eliding large resource strings in PassManager #149187

Uh oh!

akshaykhadse commented Jul 16, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 16, 2025 •

edited

Loading

Uh oh!

llvmbot commented Jul 16, 2025

Uh oh!

Uh oh!

makslevental left a comment •

edited

Loading

Uh oh!

github-actions bot commented Jul 17, 2025 •

edited

Loading

Uh oh!

makslevental commented Jul 17, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[MLIR][Python] Support eliding large resource strings in PassManager #149187

[MLIR][Python] Support eliding large resource strings in PassManager #149187

Uh oh!

Conversation

akshaykhadse commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jul 16, 2025

Uh oh!

Uh oh!

makslevental left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

makslevental commented Jul 17, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

akshaykhadse commented Jul 16, 2025 •

edited

Loading

github-actions bot commented Jul 16, 2025 •

edited

Loading

makslevental left a comment •

edited

Loading

github-actions bot commented Jul 17, 2025 •

edited

Loading